import numpy as np
from scipy.fftpack import dct, idct, fft, ifft
import matplotlib.pyplot as plt
import pandas as pd

# 参数设置
Nx, Ny = 32, 32  # 假设我们的矩阵大小为 8x8
Lx, Ly = 1, 1  # 假设我们的域长度为 1x1
dx, dy = Lx / (Nx - 1), Ly / (Ny - 1)  # 计算网格间距

# 波数
kx = 2  # 假设我们在 x 方向有一个特定的波数 kx
ky = 2  # 假设我们在 y 方向有一个特定的波数 ky

# 创建一个二维余弦函数作为输入
x = np.linspace(0, Lx, Nx, endpoint=False)
y = np.linspace(0, Ly, Ny, endpoint=False)
X, Y = np.meshgrid(x, y, indexing='ij')
u = np.cos(2 * np.pi * kx * X) * np.cos(2 * np.pi * ky * Y)

# 对每一行执行DCT
u_dct_row = dct(u, axis=1, norm='ortho')    #后指标

# 对每一列执行DCT
#u_dct_col = dct(u.T, axis=1, norm='ortho').T    #前指标
u_dct_col = dct(u.T, axis=0, norm='ortho')    #前指标

# 验证：使用IDCT重建原始矩阵
u_reconstructed_row = idct(u_dct_row, axis=1, norm='ortho')
u_reconstructed_col = idct(u_dct_col, axis=0, norm='ortho')

# 计算误差
error_row = np.linalg.norm(u - u_reconstructed_row, 'fro')
error_col = np.linalg.norm(u - u_reconstructed_col, 'fro')

print("Error after IDCT on rows:", error_row)
print("Error after IDCT on columns:", error_col)

# 检查DCT的峰值位置
print("DCT row-wise peak at:", np.unravel_index(np.argmax(u_dct_row, axis=None), u_dct_row.shape))
print("DCT column-wise peak at:", np.unravel_index(np.argmax(u_dct_col, axis=None), u_dct_col.shape))

# data = pd.DataFrame(u_dct_col)
# writer = pd.ExcelWriter(f'u_dct_col.xlsx')		# 写入Excel文件
# data.to_excel(writer, 'page_1', float_format='%.6f')		# ‘page_1’是写入excel的sheet名
# writer.close()
print(u_dct_col[4,:])   #对前指标DCT了

plt.figure(figsize=(12, 6))

plt.subplot(1, 2, 1)
plt.title("Numerical Solution")
plt.imshow(np.real(u), extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')

plt.subplot(1, 2, 2)
plt.title("Analytical Solution")
plt.imshow(np.real(u_dct_col), extent=[-1, 1, -1, 1], origin='lower', cmap='viridis')
plt.colorbar(label='Pressure')


plt.tight_layout()
plt.show()